﻿using System.Collections.Generic;
using System.Linq;
using System.Reflection.Emit;
using System.Reflection;

namespace Axial
{
	public static class ILOptimizer
	{
		// flip the previous condition and target
		private static Dictionary<OpCode, OpCode> OppositeCodes = new Dictionary<OpCode, OpCode> {
		   { OpCodes.Beq, OpCodes.Bne_Un },
		   { OpCodes.Beq_S, OpCodes.Bne_Un_S },
		   { OpCodes.Bge, OpCodes.Blt },
		   { OpCodes.Bge_S, OpCodes.Blt_S },
		   { OpCodes.Bge_Un, OpCodes.Blt_Un },
		   { OpCodes.Bge_Un_S, OpCodes.Blt_Un_S },
		   { OpCodes.Bgt, OpCodes.Ble },
		   { OpCodes.Bgt_S, OpCodes.Ble_S },
		   { OpCodes.Bgt_Un, OpCodes.Ble_Un },
		   { OpCodes.Bgt_Un_S, OpCodes.Ble_Un_S },
		   { OpCodes.Ble, OpCodes.Bgt },
		   { OpCodes.Ble_S, OpCodes.Bgt_S },
		   { OpCodes.Ble_Un, OpCodes.Bgt_Un },
		   { OpCodes.Ble_Un_S, OpCodes.Bgt_Un_S },
		   { OpCodes.Blt, OpCodes.Bge },
		   { OpCodes.Blt_S, OpCodes.Bge_S },
		   { OpCodes.Blt_Un, OpCodes.Bge_Un },
		   { OpCodes.Blt_Un_S, OpCodes.Bge_Un_S },
		   { OpCodes.Bne_Un, OpCodes.Beq },
		   { OpCodes.Bne_Un_S, OpCodes.Beq_S }
		};

		public static List<ILInstruction> Optimize(List<ILInstruction> instructions)
		{
			List<ILInstruction> optimized = new List<ILInstruction>(instructions);

			// nops are confusing
			RemoveNops(optimized);
			ChangeGetToLdfld(optimized);
			FixDoubleCompare(optimized);
			FixCompareThenBranchTrue(optimized);
			FixExtendedReturn(optimized);
			FixCompareThenBranchInWhile(optimized);
			FixTwoStepBranchInWhile(optimized);
			FixConditionalBranchThenUnconditionalBranch(optimized);
			FixPushPop(optimized);

			return optimized;
		}

		private static void FixPushPop(List<ILInstruction> optimized)
		{
			for (int i = 0; i < optimized.Count - 1; ++i)
			{
				if (optimized[i].IsLoadCode() && optimized[i + 1].Code == OpCodes.Pop)
				{
					SaveRemoveInstructionRange(optimized, i, 2);
					i -= 2;
				}
			}
		}

		private static void FixTwoStepBranchInWhile(List<ILInstruction> optimized)
		{
			// check for vb.net debug mode deoptimization
			foreach (ILInstruction branch in optimized.FindAll(il => il.IsUnconditionalBranch()))
			{
				// match the pattern
				int branchOffset = optimized.IndexOf(branch);
				ILInstruction loadZero = optimized[branchOffset - 1];
				if (loadZero.Code != OpCodes.Ldc_I4_0)
					continue;
				if (branch.GetBranchTarget() != optimized[branchOffset + 2].Offset)
					continue;
				ILInstruction loadOne = optimized[branchOffset + 1];
				if (loadOne.Code != OpCodes.Ldc_I4_1)
					continue;
				ILInstruction store = optimized[branchOffset + 2];
				ILInstruction load = optimized[branchOffset + 3];
				if (!store.IsStoreCode() || !load.IsLoadCode() || store.GetStoreLocation() != load.GetLoadLocation())
					continue;
				ILInstruction lastbr = optimized[branchOffset + 4];
				if (!lastbr.IsBrTrue() && !lastbr.IsBrFalse())
					continue;

				int trueCodeOffset = optimized[branchOffset + 5].Offset;
				if (lastbr.IsBrFalse())
				{
					// change all branches targeting loadOne
					foreach (ILInstruction inst in optimized.Take(branchOffset - 1).Where(il => il.IsBranch() && il.GetBranchTarget() == loadOne.Offset))
						inst.Operand = trueCodeOffset;

					// change all branches targeting loadZero
					foreach (ILInstruction inst in optimized.Take(branchOffset - 1).Where(il => il.IsBranch() && il.GetBranchTarget() == loadZero.Offset))
						inst.Operand = lastbr.Operand;
				}
				else
				{
					// change all branches targeting loadOne
					foreach (ILInstruction inst in optimized.Take(branchOffset - 1).Where(il => il.IsBranch() && il.GetBranchTarget() == loadOne.Offset))
						inst.Operand = lastbr.Operand;

					// change all branches targeting loadZero
					foreach (ILInstruction inst in optimized.Take(branchOffset - 1).Where(il => il.IsBranch() && il.GetBranchTarget() == loadZero.Offset))
						inst.Operand = trueCodeOffset;
				}

				ILInstruction lastConditional = optimized[branchOffset - 2];
				if (lastConditional.GetBranchTarget() == trueCodeOffset)
					lastConditional.Operand = lastbr.Operand;
				else
					lastConditional.Operand = trueCodeOffset;
				lastConditional.Code = OppositeCodes[lastConditional.Code];

				// remove the useless instructions
				SaveRemoveInstructionRange(optimized, branchOffset - 1, 6);
			}
		}

		private static void FixConditionalBranchThenUnconditionalBranch(List<ILInstruction> optimized)
		{
			// check for a different vb.net debug mode deoptimization
			foreach (ILInstruction condBranch in optimized.FindAll(il => il.IsConditionalBranch()))
			{
				// match the pattern
				int condIndex = optimized.IndexOf(condBranch);
				ILInstruction uncondBranch = optimized[condIndex + 1];
				if (!uncondBranch.IsUnconditionalBranch())
					continue;
				// i don't understand this part of the pattern but it's proven to be necessary
				if (!optimized.Take(condIndex).Any(il => il.IsBranch() &&
						(il.GetBranchTarget() == condBranch.GetBranchTarget() || il.GetBranchTarget() == uncondBranch.Offset)))
					continue;

				// flip condition and give it unconditional's target, then remove the unconditional
				condBranch.Code = OppositeCodes[condBranch.Code];
				condBranch.Operand = uncondBranch.Operand;
				foreach (ILInstruction inst in optimized.Take(condIndex).Where(il => il.IsConditionalBranch() && il.GetBranchTarget() == uncondBranch.Offset))
					inst.Operand = uncondBranch.Operand;
				SaveRemoveInstruction(optimized, condIndex + 1);
			}
		}

		private static void FixCompareThenBranchInWhile(List<ILInstruction> optimized)
		{
			// check for c# debug mode while oddity
			foreach (ILInstruction comp in optimized.FindAll(il => il != optimized.Last() && il.IsComparison() &&
				optimized[optimized.IndexOf(il) + 1].IsUnconditionalBranch()))
			{
				int index = optimized.IndexOf(comp);
				ILInstruction skip = optimized[index + 1];
				ILInstruction overwrite = optimized[index + 2];

				// find store, load and brtrue instructions
				ILInstruction store = optimized.Skip(index).First(il => il.IsStoreCode());
				int storeindex = optimized.IndexOf(store);
				ILInstruction load = optimized[storeindex + 1];
				ILInstruction end = optimized[storeindex + 2];

				int oldtarget = overwrite.Offset;
				int trueCodeOffset = (int)end.Operand;
				int postLoopOffset = optimized[storeindex + 3].Offset;
				int newtarget = overwrite.Code == OpCodes.Ldc_I4_1 ? trueCodeOffset : postLoopOffset;
				// change branches pointing to the overwrite
				foreach (ILInstruction inst in optimized.Take(index).Where(il => il.IsBranch() && (int)il.Operand == oldtarget))
					inst.Operand = newtarget;

				// change the comparison
				comp.Operand = trueCodeOffset;
				if (comp.Code == OpCodes.Ceq)
					comp.Code = OpCodes.Beq_S;
				if (comp.Code == OpCodes.Cgt)
					comp.Code = OpCodes.Bgt_S;
				if (comp.Code == OpCodes.Clt)
					comp.Code = OpCodes.Blt_S;

				// see if there was both a true and false shortcut
				if (storeindex - index != 3)
				{
					overwrite = optimized[index + 4];
					oldtarget = overwrite.Offset;
					newtarget = overwrite.Code == OpCodes.Ldc_I4_1 ? trueCodeOffset : postLoopOffset;
					// change branches pointing to the overwrite
					foreach (ILInstruction inst in optimized.Take(index).Where(il => il.IsBranch() && (int)il.Operand == oldtarget))
						inst.Operand = newtarget;
					SaveRemoveInstructionRange(optimized, index + 1, 2);
				}

				SaveRemoveInstructionRange(optimized, index + 1, 5);
			}
		}

		private static void FixExtendedReturn(List<ILInstruction> optimized)
		{
			// check for an extended return: store, branch to next instruction, load what was stored, return
			if (optimized.Count > 4)
			{
				var ret = optimized[optimized.Count - 1];
				var load = optimized[optimized.Count - 2];
				var branch = optimized[optimized.Count - 3];
				var store = optimized[optimized.Count - 4];

				if (ret.Code == OpCodes.Ret &&
					store.IsStoreCode() && load.IsLoadCode() && store.GetStoreLocation() == load.GetLoadLocation() &&
					branch.Code.FlowControl == FlowControl.Branch && (int)branch.Operand == branch.Offset + 2)
					SaveRemoveInstructionRange(optimized, optimized.Count - 4, 3);
			}
		}

		private static void FixCompareThenBranchTrue(List<ILInstruction> optimized)
		{
			// check for alternate odd conditional: ceq/clt/cgt, store, load, brtrue
			if (optimized.Count > 4)
			{
				// c# debug mode uses brtrue opcode
				Dictionary<OpCode, OpCode> truematches = new Dictionary<OpCode, OpCode> {
					{ OpCodes.Ceq, OpCodes.Beq_S }, 
					{ OpCodes.Cgt, OpCodes.Bgt_Un_S }, 
					{ OpCodes.Clt, OpCodes.Blt_Un_S } 
				};
				// vb.net debug mode uses brfalse opcodes
				Dictionary<OpCode, OpCode> falsematches = new Dictionary<OpCode, OpCode> {
					{ OpCodes.Ceq, OpCodes.Bne_Un_S }, 
					{ OpCodes.Cgt, OpCodes.Ble_Un_S }, 
					{ OpCodes.Clt, OpCodes.Bge_Un_S } 
				};
				foreach (ILInstruction comp in optimized.FindAll(il => truematches.Keys.Contains(il.Code)))
				{
					int search = optimized.IndexOf(comp);
					if (search > optimized.Count - 4)
						break;

					ILInstruction store = optimized[search + 1];
					ILInstruction load = optimized[search + 2];
					ILInstruction brtrue = optimized[search + 3];

					// make sure the pattern matches
					if (!store.IsStoreCode() || !load.IsLoadCode() || store.GetStoreLocation() != load.GetLoadLocation())
						continue;
					// c# check
					if (brtrue.IsBrTrue())
					{
						brtrue.Code = truematches[comp.Code];
						SaveRemoveInstructionRange(optimized, search, 3);
					}
					// vb check
					if (brtrue.IsBrFalse())
					{
						brtrue.Code = falsematches[comp.Code];
						SaveRemoveInstructionRange(optimized, search, 3);
					}
				}
			}
		}

		private static void FixDoubleCompare(List<ILInstruction> optimized)
		{
			// look for c# debug mode oddity with branching shortcuts
			if (optimized.Count > 5)
			{
				Dictionary<OpCode, OpCode> brtrueMatches = new Dictionary<OpCode, OpCode> {
					{ OpCodes.Ceq, OpCodes.Bne_Un_S }, 
					{ OpCodes.Cgt, OpCodes.Ble_Un_S },
					{ OpCodes.Clt, OpCodes.Bge_Un_S },
				};
				Dictionary<OpCode, OpCode> brfalseMatches = new Dictionary<OpCode, OpCode> {
					{ OpCodes.Ceq, OpCodes.Beq_S }, 
					{ OpCodes.Cgt, OpCodes.Bgt_Un_S },
					{ OpCodes.Clt, OpCodes.Blt_Un_S },
				};
				foreach (ILInstruction comp in optimized.FindAll(il => brtrueMatches.Keys.Contains(il.Code)))
				{
					int search = optimized.IndexOf(comp);
					// make sure the pattern matches
					if (optimized[search + 2].Code != OpCodes.Ceq)
						continue;

					// remove ldc_i4_0 and extra ceq
					SaveRemoveInstructionRange(optimized, search + 1, 2);
					// if there's a true shortcut, modify all shortcuts that lead to it to go to the true code
					if (optimized[search + 2].Code == OpCodes.Ldc_I4_1)
					{
						foreach (var newJump in optimized.Where(il => il.IsBranch() && (int)il.Operand == optimized[search + 2].Offset))
							newJump.Operand = optimized.Skip(search + 2).First(il => il.IsBranch()).Operand;
						if (optimized[search + 1].IsUnconditionalBranch())
							SaveRemoveInstructionRange(optimized, search, 2);
					}
					// do the same: change all false shortcuts to lead to false code
					if (optimized[search + 2].Code == OpCodes.Ldc_I4_0)
					{
						foreach (var newJump in optimized.Where(il => il.IsBranch() && (int)il.Operand == optimized[search + 2].Offset))
							newJump.Operand = optimized[optimized.IndexOf(optimized.Skip(search + 2).First(il => il.IsBranch())) + 1].Offset;
						if (optimized[search + 1].IsUnconditionalBranch())
							SaveRemoveInstructionRange(optimized, search, 2);
					}

					// remove the bad condition
					SaveRemoveInstruction(optimized, search);
					// remove the store and load
					if (optimized[search].IsStoreCode() && optimized[search + 1].IsLoadCode() &&
						optimized[search].GetStoreLocation() == optimized[search + 1].GetLoadLocation())
						SaveRemoveInstructionRange(optimized, search, 2);

					// fix the opcode so the branch has the correct condition
					// c# debug compiler puts brtrue, vb debug compiler uses brfalse
					var matchDict = optimized[search].IsBrTrue() ? brtrueMatches : brfalseMatches;
					optimized[search].Code = matchDict[comp.Code];
				}
			}
		}

		private static void ChangeGetToLdfld(List<ILInstruction> optimized)
		{
			// fix for VB.NET compiler - change get_ methods to ldfld where appropriate
			foreach (ILInstruction inst in optimized.Where(il => il.Code == OpCodes.Call || il.Code == OpCodes.Callvirt))
			{
				MethodBase method = (MethodBase)inst.Operand;
				if (!method.Name.StartsWith("get_"))
					continue;
				FieldInfo possibleField = method.DeclaringType.GetField(method.Name.Substring(3), BindingFlags.NonPublic | BindingFlags.Instance);
				if (possibleField != null)
				{
					inst.Code = OpCodes.Ldfld;
					inst.Operand = possibleField;
				}
			}
		}

		// make sure anything that branched to a nop branches to the next code
		private static void RemoveNops(List<ILInstruction> optimized)
		{
			foreach (ILInstruction nop in optimized.FindAll(il => il.Code == OpCodes.Nop))
			{
				SaveRemoveInstruction(optimized, nop);
			}
		}

		private static void SaveRemoveInstruction(List<ILInstruction> optimized, int index)
		{
			ILInstruction toRemove = optimized[index];
			optimized.RemoveAt(index);
			foreach (ILInstruction after in optimized.Skip(index))
				after.Offset -= 1;
			foreach (ILInstruction after in optimized.FindAll(o => o.GetBranchTarget() > toRemove.Offset))
				after.Operand = (int)after.Operand - 1;
		}

		private static void SaveRemoveInstructionRange(List<ILInstruction> optimized, int index, int count)
		{
			foreach (var i in Enumerable.Range(0, count))
				SaveRemoveInstruction(optimized, index);
		}

		private static void SaveRemoveInstruction(List<ILInstruction> optimized, ILInstruction toRemove)
		{
			SaveRemoveInstruction(optimized, optimized.IndexOf(toRemove));
		}
	}
}
